#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Curve Fitting Tool
Function: Try multiple fitting methods on two numeric data columns and select the best fitting method and curve

Version: 1.0.0
Author: Assistant
"""

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import argparse
import sys
from pathlib import Path
from scipy import stats
from scipy.optimize import curve_fit
from sklearn.metrics import r2_score, mean_squared_error
from sklearn.preprocessing import PolynomialFeatures
from sklearn.linear_model import LinearRegression
from sklearn.pipeline import Pipeline
import warnings
import os

# Set matplotlib font support
plt.rcParams['font.sans-serif'] = ['DejaVu Sans', 'Arial', 'sans-serif']
plt.rcParams['axes.unicode_minus'] = False


class CurveFitter:
    """
    Curve Fitting Tool Class
    
    Function:
    - Support reading CSV/Excel format data files
    - Try multiple fitting methods on two numeric data columns
    - Select the best fitting method
    - Generate scatter plots for each fitting method, including formulas and necessary information
    """
    
    def __init__(self):
        """
        Initialize the curve fitting tool
        """
        self.data = None
        self.col1 = None
        self.col2 = None
        self.output_dir = None
        self.prefix = None
        self.fitted_models = {}
        self.best_model = None
        self.best_model_name = None
        self.best_r2 = -np.inf
        
        # Define fitting functions
        self.fit_functions = {
            'linear': self._linear_func,
            'quadratic': self._quadratic_func,
            'cubic': self._cubic_func,
            'exponential': self._exponential_func,
            'logarithmic': self._logarithmic_func,
            'power': self._power_func
        }
    
    # Define fitting functions
    def _linear_func(self, x, a, b):
        """Linear function: y = a*x + b"""
        return a * x + b
    
    def _quadratic_func(self, x, a, b, c):
        """Quadratic function: y = a*x^2 + b*x + c"""
        return a * x**2 + b * x + c
    
    def _cubic_func(self, x, a, b, c, d):
        """Cubic function: y = a*x^3 + b*x^2 + c*x + d"""
        return a * x**3 + b * x**2 + c * x + d
    
    def _exponential_func(self, x, a, b, c):
        """Exponential function: y = a * exp(b*x) + c"""
        return a * np.exp(b * x) + c
    
    def _logarithmic_func(self, x, a, b):
        """Logarithmic function: y = a * ln(x) + b"""
        # Handle x<=0 cases
        x_safe = np.where(x > 0, x, 1e-10)
        return a * np.log(x_safe) + b
    
    def _power_func(self, x, a, b):
        """Power function: y = a * x^b"""
        # Handle x<=0 cases
        x_safe = np.where(x > 0, x, 1e-10)
        return a * np.power(x_safe, b)
    
    def load_data(self, file_path):
        """
        Load data file
        
        Args:
            file_path (str): Data file path
            
        Returns:
            bool: Returns True if loading is successful, otherwise False
        """
        try:
            file_path = Path(file_path)
            
            if not file_path.exists():
                print(f"Error: File {file_path} does not exist")
                return False
            
            # Select reading method based on file extension
            if file_path.suffix.lower() == '.csv':
                self.data = pd.read_csv(file_path)
            elif file_path.suffix.lower() in ['.xlsx', '.xls']:
                self.data = pd.read_excel(file_path)
            else:
                print(f"Error: Unsupported file format {file_path.suffix}")
                return False
            
            print(f"Successfully loaded data: {len(self.data)} rows, {len(self.data.columns)} columns")
            print(f"Column names: {list(self.data.columns)}")
            return True
            
        except Exception as e:
            print(f"Error occurred while loading data: {str(e)}")
            return False
    
    def validate_columns(self, col1, col2):
        """
        Validate that column names exist and are numeric
        
        Args:
            col1 (str): First column name
            col2 (str): Second column name
            
        Returns:
            bool: Returns True if validation passes, otherwise False
        """
        if self.data is None:
            print("Error: Please load data first")
            return False
        
        # Check if columns exist
        if col1 not in self.data.columns:
            print(f"Error: Column '{col1}' does not exist")
            return False
        
        if col2 not in self.data.columns:
            print(f"Error: Column '{col2}' does not exist")
            return False
        
        # Check if columns are numeric
        if not pd.api.types.is_numeric_dtype(self.data[col1]):
            print(f"Error: Column '{col1}' is not numeric")
            return False
        
        if not pd.api.types.is_numeric_dtype(self.data[col2]):
            print(f"Error: Column '{col2}' is not numeric")
            return False
        
        # Set column names
        self.col1 = col1
        self.col2 = col2
        
        return True
    
    def set_output_config(self, output_dir, prefix):
        """
        Set output configuration
        
        Args:
            output_dir (str): Output directory
            prefix (str): Output file prefix
        """
        self.output_dir = Path(output_dir)
        self.prefix = prefix
        
        # Ensure output directory exists
        self.output_dir.mkdir(parents=True, exist_ok=True)
    
    def _fit_model(self, func_name, func, x_data, y_data):
        """
        Fit a single model
        
        Args:
            func_name (str): Function name
            func (callable): Fitting function
            x_data (np.array): X data
            y_data (np.array): Y data
            
        Returns:
            dict: Fitting results
        """
        try:
            # Determine initial parameters based on function parameter count
            if func_name == 'linear':
                p0 = [1.0, 1.0]
            elif func_name == 'quadratic':
                p0 = [1.0, 1.0, 1.0]
            elif func_name == 'cubic':
                p0 = [1.0, 1.0, 1.0, 1.0]
            elif func_name == 'exponential':
                p0 = [1.0, 0.1, 1.0]
            elif func_name == 'logarithmic':
                p0 = [1.0, 1.0]
            elif func_name == 'power':
                p0 = [1.0, 1.0]
            else:
                p0 = None
            
            # Perform fitting
            popt, pcov = curve_fit(func, x_data, y_data, p0=p0, maxfev=10000)
            
            # Calculate predicted values
            y_pred = func(x_data, *popt)
            
            # Calculate evaluation metrics
            r2 = r2_score(y_data, y_pred)
            rmse = np.sqrt(mean_squared_error(y_data, y_pred))
            
            # Calculate parameter standard errors
            param_errors = np.sqrt(np.diag(pcov)) if pcov.size > 0 else [0] * len(popt)
            
            return {
                'params': popt,
                'param_errors': param_errors,
                'r2': r2,
                'rmse': rmse,
                'y_pred': y_pred,
                'function': func
            }
        except Exception as e:
            print(f"Error fitting {func_name} model: {str(e)}")
            return None
    
    def fit_all_models(self):
        """
        Fit all models
        """
        if self.data is None or self.col1 is None or self.col2 is None:
            raise ValueError("Please load data and set column names first")
        
        # Get data
        x_data = self.data[self.col1].values
        y_data = self.data[self.col2].values
        
        # Remove NaN values
        mask = ~(np.isnan(x_data) | np.isnan(y_data))
        x_clean = x_data[mask]
        y_clean = y_data[mask]
        
        if len(x_clean) < 3:
            raise ValueError("Less than 3 valid data points, unable to perform fitting")
        
        print(f"Starting to fit {len(self.fit_functions)} models...")
        
        # Try each fitting method
        for name, func in self.fit_functions.items():
            print(f"Fitting {name} model...")
            result = self._fit_model(name, func, x_clean, y_clean)
            if result is not None:
                self.fitted_models[name] = result
                print(f"  R² = {result['r2']:.4f}, RMSE = {result['rmse']:.4f}")
                
                # Update best model
                if result['r2'] > self.best_r2:
                    self.best_r2 = result['r2']
                    self.best_model = result
                    self.best_model_name = name
            else:
                print(f"  Fitting failed")
        
        if self.best_model is not None:
            print(f"\nBest fitting model: {self.best_model_name} (R² = {self.best_r2:.4f})")
        else:
            print("\nAll model fittings failed")
    
    def _format_equation(self, model_name, params):
        """
        Format equation display
        
        Args:
            model_name (str): Model name
            params (list): Parameter list
            
        Returns:
            str: Formatted equation
        """
        if model_name == 'linear':
            a, b = params
            return f"y = {a:.4f}*x + {b:.4f}"
        elif model_name == 'quadratic':
            a, b, c = params
            return f"y = {a:.4f}*x² + {b:.4f}*x + {c:.4f}"
        elif model_name == 'cubic':
            a, b, c, d = params
            return f"y = {a:.4f}*x³ + {b:.4f}*x² + {c:.4f}*x + {d:.4f}"
        elif model_name == 'exponential':
            a, b, c = params
            return f"y = {a:.4f}*exp({b:.4f}*x) + {c:.4f}"
        elif model_name == 'logarithmic':
            a, b = params
            return f"y = {a:.4f}*ln(x) + {b:.4f}"
        elif model_name == 'power':
            a, b = params
            return f"y = {a:.4f}*x^{b:.4f}"
        else:
            return "Unknown model"
    
    def generate_plots(self):
        """
        Generate charts for all fitted models
        """
        if not self.fitted_models:
            print("No fitted models to plot")
            return
        
        # Get data
        x_data = self.data[self.col1].values
        y_data = self.data[self.col2].values
        
        # Remove NaN values
        mask = ~(np.isnan(x_data) | np.isnan(y_data))
        x_clean = x_data[mask]
        y_clean = y_data[mask]
        
        # Generate charts for each model
        for name, model in self.fitted_models.items():
            # Create chart
            fig, ax = plt.subplots(figsize=(10, 8))
            
            # Plot original data scatter
            ax.scatter(x_clean, y_clean, alpha=0.7, color='blue', s=50, label='Original Data')
            
            # Generate fitting curve
            x_fit = np.linspace(x_clean.min(), x_clean.max(), 200)
            y_fit = model['function'](x_fit, *model['params'])
            
            # Plot fitting curve
            ax.plot(x_fit, y_fit, 'r-', linewidth=2, label=f'{name} Fit')
            
            # Add equation and statistical information
            equation = self._format_equation(name, model['params'])
            r2 = model['r2']
            rmse = model['rmse']
            
            # Add information to chart
            info_text = f"Equation: {equation}\nR² = {r2:.4f}\nRMSE = {rmse:.4f}"
            ax.text(0.05, 0.95, info_text, transform=ax.transAxes, 
                   verticalalignment='top', bbox=dict(boxstyle='round', facecolor='white', alpha=0.8),
                   fontsize=10, fontfamily='monospace')
            
            # Set chart properties
            ax.set_xlabel(self.col1, fontsize=12)
            ax.set_ylabel(self.col2, fontsize=12)
            ax.set_title(f'{name.capitalize()} Fitting Curve', fontsize=14, fontweight='bold')
            ax.grid(True, alpha=0.3)
            ax.legend()
            
            # Save chart
            filename = f"{self.prefix}_{name}_fit.png"
            filepath = self.output_dir / filename
            plt.savefig(filepath, dpi=300, bbox_inches='tight', facecolor='white')
            plt.close(fig)
            
            print(f"Saved {name} fitting chart: {filepath}")
        
        # Generate comparison chart for best fitting model
        if self.best_model is not None:
            fig, ax = plt.subplots(figsize=(12, 8))
            
            # Plot original data scatter
            ax.scatter(x_clean, y_clean, alpha=0.7, color='blue', s=50, label='Original Data')
            
            # Generate best fitting curve
            x_fit = np.linspace(x_clean.min(), x_clean.max(), 200)
            y_fit = self.best_model['function'](x_fit, *self.best_model['params'])
            
            # Plot best fitting curve
            ax.plot(x_fit, y_fit, 'r-', linewidth=2, label=f'Best Fit ({self.best_model_name})')
            
            # Add equation and statistical information
            equation = self._format_equation(self.best_model_name, self.best_model['params'])
            r2 = self.best_model['r2']
            rmse = self.best_model['rmse']
            
            # Add information to chart
            info_text = f"Best Fit Model: {self.best_model_name}\nEquation: {equation}\nR² = {r2:.4f}\nRMSE = {rmse:.4f}"
            ax.text(0.05, 0.95, info_text, transform=ax.transAxes, 
                   verticalalignment='top', bbox=dict(boxstyle='round', facecolor='lightgreen', alpha=0.8),
                   fontsize=11, fontfamily='monospace')
            
            # Set chart properties
            ax.set_xlabel(self.col1, fontsize=12)
            ax.set_ylabel(self.col2, fontsize=12)
            ax.set_title(f'Best Fitting Curve Comparison', fontsize=14, fontweight='bold')
            ax.grid(True, alpha=0.3)
            ax.legend()
            
            # Save chart
            filename = f"{self.prefix}_best_fit.png"
            filepath = self.output_dir / filename
            plt.savefig(filepath, dpi=300, bbox_inches='tight', facecolor='white')
            plt.close(fig)
            
            print(f"Saved best fit comparison chart: {filepath}")
    
    def print_summary(self):
        """
        Print fitting result summary
        """
        if not self.fitted_models:
            print("No fitting results to display")
            return
        
        print("\n=== Fitting Result Summary ===")
        print(f"Data columns: {self.col1} vs {self.col2}")
        print(f"Number of data points: {len(self.data)}")
        
        print("\nResults for each model:")
        print("-" * 80)
        print(f"{'Model':<12} {'R²':<10} {'RMSE':<10} {'Best':<8}")
        print("-" * 80)
        
        for name, model in self.fitted_models.items():
            r2 = model['r2']
            rmse = model['rmse']
            is_best = "Yes" if name == self.best_model_name else "No"
            print(f"{name:<12} {r2:<10.4f} {rmse:<10.4f} {is_best:<8}")
        
        if self.best_model_name:
            print(f"\nBest fitting model: {self.best_model_name}")
            print(f"Best R² value: {self.best_r2:.4f}")
            equation = self._format_equation(self.best_model_name, self.best_model['params'])
            print(f"Fitting equation: {equation}")


def main():
    """
    Main function - Command line interface
    """
    parser = argparse.ArgumentParser(
        description='Curve Fitting Tool - Try multiple fitting methods on two numeric data columns',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
Usage Examples:
  # Basic usage
  python curve_fitter.py -i data.csv --c1 column1 --c2 column2 -o ./output -p result
  
  # Using Excel file
  python curve_fitter.py -i data.xlsx --c1 x_data --c2 y_data -o ./fits -p experiment1
        """
    )
    
    # Input parameters
    parser.add_argument('-i', '--input', required=True,
                       help='Input data file path (CSV or Excel format)')
    parser.add_argument('--c1', required=True,
                       help='First column name (must be numeric)')
    parser.add_argument('--c2', required=True,
                       help='Second column name (must be numeric)')
    
    # Output parameters
    parser.add_argument('-o', '--output', required=True,
                       help='Output directory path')
    parser.add_argument('-p', '--prefix', required=True,
                       help='Output file prefix')
    
    args = parser.parse_args()
    
    try:
        # Create curve fitting tool instance
        fitter = CurveFitter()
        
        # Load data
        if not fitter.load_data(args.input):
            return 1
        
        # Validate column names
        if not fitter.validate_columns(args.c1, args.c2):
            return 1
        
        # Set output configuration
        fitter.set_output_config(args.output, args.prefix)
        
        # Perform fitting for all models
        fitter.fit_all_models()
        
        # Print result summary
        fitter.print_summary()
        
        # Generate charts
        print("\nGenerating fitting charts...")
        fitter.generate_plots()
        
        print(f"\nFitting completed! Results saved to directory: {args.output}")
        return 0
        
    except KeyboardInterrupt:
        print("\nUser interrupted operation")
        return 1
    except Exception as e:
        print(f"Error occurred during program execution: {str(e)}")
        return 1


if __name__ == '__main__':
    exit(main())